package com.aleyn.router.plug.task

import com.aleyn.router.plug.RouterPlugin
import com.aleyn.router.plug.data.HandleModel
import com.aleyn.router.plug.visitor.FindHandleClass
import com.aleyn.router.plug.visitor.InsertCodeVisitor
import com.android.build.gradle.internal.cxx.io.writeTextIfDifferent
import org.gradle.api.DefaultTask
import org.gradle.api.Project
import org.gradle.api.file.Directory
import org.gradle.api.file.DirectoryProperty
import org.gradle.api.file.RegularFile
import org.gradle.api.file.RegularFileProperty
import org.gradle.api.provider.ListProperty
import org.gradle.api.provider.Property
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.InputFile
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.OutputDirectory
import org.gradle.api.tasks.OutputFile
import org.gradle.api.tasks.TaskAction
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter
import java.io.BufferedOutputStream
import java.io.File
import java.io.FileOutputStream
import java.io.InputStream
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream

/**
 * @author: Aleyn
 * @date: 2023/7/14 11:03
 *
 * Gradle 任务，用于扫描和处理 LRouter 的注解。
 *
 * 这个任务是 LRouter 框架实现自动注册的核心。它的主要职责包括：
 * 1. **扫描字节码**: 遍历项目编译后的所有 .class 文件和依赖的 .jar 包。
 * 2. **识别注解**: 使用 ASM 技术识别被 `@Route`, `@Interceptor`, `@Initializer`, `@Autowired` 等注解标记的类。
 * 3. **信息缓存**: 将扫描到的元信息（如类名、路径、优先级等）分类存储到构建目录下的临时文件中（例如 `build/intermediates/router_scanning_data/`）。
 * 4. **代码注入**: 在主应用模块（app module）中，找到预定义的桩类 `LRouterGenerateImpl`，并通过 ASM 将所有模块缓存的元信息读取并注册到该类中。
 *
 * 通过这种方式，LRouter 避免了在运行时使用反射来扫描类，从而显著提高了应用的启动性能和路由表的加载速度。
 */

/**
 * 定义了待插桩的目标类。
 * 这是由 `RouterStubClassTask` 任务预先生成的一个桩类，包含了空的方法实现。
 * `LRouterClassTask` 的任务就是向这个类中填充具体的路由注册逻辑。
 */
internal const val ROUTER_GENERATE = "com/router/LRouterGenerateImpl.class"

/**
 * LRouter 字节码扫描与代码注入任务。
 *
 * 继承自 Gradle 的 `DefaultTask`。
 * 通过 `@InputFiles`, `@OutputFile`, `@Input` 等注解声明任务的输入和输出，以支持 Gradle 的增量构建。
 */
abstract class LRouterClassTask : DefaultTask() {

    /**
     * 输入：所有依赖的 JAR 包。
     * `ListProperty<RegularFile>` 用于接收一个文件列表。
     */
    @get:InputFiles
    abstract val allJars: ListProperty<RegularFile>

    /**
     * 输入：所有编译后的 class 文件所在的目录。
     * 通常是 `build/intermediates/javac/debug/classes` 等。
     */
    @get:InputFiles
    abstract val allDirectories: ListProperty<Directory>

    /**
     * 输出：处理后的 JAR 文件。
     * 这个 JAR 文件包含了所有原始的类以及被修改后的 `LRouterGenerateImpl` 类。
     */
    @get:OutputFile
    abstract val output: RegularFileProperty

    /**
     * 输入：用于存储扫描信息的缓存路径。
     * 例如 `build/intermediates/router_scanning_data/debug`。
     */
    @get:Input
    abstract val cachePath: Property<String>

    /**
     * Gradle 任务的核心执行逻辑。
     *
     * 此方法由 Gradle 在构建过程中自动调用。其工作流程如下：
     * 1.  **初始化**: 创建一个 `JarOutputStream` 用于写入最终的输出 JAR 文件。初始化一个列表 `handleModels` 来收集所有扫描到的注解信息。
     * 2.  **扫描目录**: 遍历 `allDirectories` 中的所有 `.class` 文件。
     *     - 对每个文件，调用 `findClass` 方法进行 ASM 扫描，将结果存入 `handleModels`。
     *     - 将文件复制到输出的 JAR 包中。
     *     - 如果遇到 `ROUTER_GENERATE`（即 `LRouterGenerateImpl.class`），则将其暂存起来（`waitInsertJar`），因为它需要最后被修改。
     * 3.  **扫描 JAR 包**: 遍历 `allJars` 中的所有 JAR 文件。
     *     - 对 JAR 中的每个条目（`JarEntry`），如果是 `.class` 文件，则同样调用 `findClass` 进行扫描。
     *     - 将所有条目（无论是否为 class）都复制到输出的 JAR 包中。
     * 4.  **写入缓存**: 调用 `write` 扩展函数，将 `handleModels` 中收集到的信息按类型（module, interceptor 等）写入到指定的缓存目录中。
     * 5.  **代码注入 (仅限 App 模块)**:
     *     - 判断当前项目是否为主应用模块 (`project.isApp()`)。
     *     - 如果是，则聚合所有子模块的缓存文件目录。
     *     - 使用 ASM 的 `ClassReader` 读取之前暂存的 `LRouterGenerateImpl.class` 的字节码。
     *     - 创建一个 `InsertCodeVisitor` (自定义的 `ClassVisitor`)，它负责向 `LRouterGenerateImpl` 的方法中插入加载并注册所有路由信息的字节码。
     *     - 将 `InsertCodeVisitor` 生成的修改后字节码写入输出的 JAR 包中，替换掉原始的 `LRouterGenerateImpl.class`。
     * 6.  **完成**: 关闭 `JarOutputStream`，完成任务。
     *
     * @see findClass
     * @see write
     * @see InsertCodeVisitor
     */
    @TaskAction
    fun taskAction() {

        // 用于存储从类中扫描到的所有注解信息
        val handleModels = arrayListOf<HandleModel>()

        // 用于暂存待插入代码的 LRouterGenerateImpl.class 文件
        var waitInsertJar: File? = null

        // 创建一个 Jar 输出流，用于生成最终的输出文件
        val jarOutput = JarOutputStream(BufferedOutputStream(FileOutputStream(output.get().asFile)))

        handleModels.clear()

        // 1. 遍历所有源码编译成的 .class 文件目录
        allDirectories.get().onEach { directory ->
            val directoryUri = directory.asFile.toURI()
            directory.asFile
                .walk()
                .filter { it.isFile }
                .forEach { file ->
                    // 计算文件在目录中的相对路径
                    val filePath = directoryUri
                        .relativize(file.toURI())
                        .path
                        .replace(File.separatorChar, '/')

                    // 如果是待插桩的类，则先保存起来，不直接写入输出流
                    if (filePath == ROUTER_GENERATE) {
                        waitInsertJar = file
                        return@forEach
                    }

                    // 将文件写入 Jar 输出流
                    jarOutput.putNextEntry(JarEntry(filePath))
                    file.inputStream().use { it.copyTo(jarOutput) }
                    jarOutput.closeEntry()

                    // 如果是 class 文件，则进行扫描
                    if (file.name.endsWith(".class")) {
                        file.inputStream().findClass(handleModels)
                    }
                }
        }

        // 2. 遍历所有依赖的 Jar 包
        allJars.get().onEach { file ->
            val jarFile = JarFile(file.asFile)
            jarFile.entries().iterator().forEach { jarEntry ->
                try {
                    // 将 JarEntry 写入新的 Jar 输出流
                    jarOutput.putNextEntry(JarEntry(jarEntry.name))
                    jarFile.getInputStream(jarEntry).use { it.copyTo(jarOutput) }

                    // 如果是 class 文件，则进行扫描
                    if (jarEntry.name.endsWith(".class")) {
                        runCatching {
                            jarFile.getInputStream(jarEntry).findClass(handleModels)
                        }.onFailure {
                            println("LRouter handle " + jarEntry.name + " error:${it.message}")
                        }
                    }
                } catch (_: Exception) {
                    // 忽略一些可能出现的错误，例如重复的 entry
                } finally {
                    jarOutput.closeEntry()
                }
            }
            jarFile.close()
        }

        // 3. 将扫描到的信息写入缓存文件
        val cacheDir = project.layout.buildDirectory.dir(cachePath.get()).get()
        handleModels.write(cacheDir)

        // 4. 如果是 App 模块，则执行代码注入
        if (project.isApp()) {
            if (waitInsertJar == null) throw RuntimeException("The class to insert was not found, please check for references LRouter")
            val cachePath = cachePath.get().substringBeforeLast("/")
            // 获取所有应用了 LRouter 插件的子模块的构建目录
            val allBuildDir = project.rootProject.subprojects
                .filter { it.plugins.hasPlugin(RouterPlugin::class.java) }
                .map { it.layout.buildDirectory.dir(cachePath).get() }

            // 开始向 LRouterGenerateImpl.class 注入代码
            jarOutput.putNextEntry(JarEntry(ROUTER_GENERATE))
            val writer = ClassWriter(ClassWriter.COMPUTE_FRAMES)
            val insertVisitor = InsertCodeVisitor(writer, allBuildDir)
            ClassReader(waitInsertJar!!.inputStream()).accept(insertVisitor, ClassReader.SKIP_DEBUG)
            // 将修改后的字节码写入 Jar 输出流
            jarOutput.write(writer.toByteArray())
            jarOutput.closeEntry()
        }
        // 关闭 Jar 输出流，完成文件写入
        jarOutput.close()
    }

}

/**
 * `InputStream` 的扩展函数，用于从 class 文件流中查找被注解标记的类。
 *
 * 它使用 ASM 的 `ClassReader` 来解析字节码，并传递一个 `FindHandleClass` 访问者。
 * `FindHandleClass` 会在访问类的各个部分（注解、方法等）时，判断是否存在 LRouter 相关的注解，
 * 如果存在，则提取信息并创建一个 `HandleModel` 对象。
 *
 * @param outHandle 用于收集扫描结果的列表。
 *
 * @see FindHandleClass
 * @see org.objectweb.asm.ClassReader.accept
 */
fun InputStream.findClass(outHandle: ArrayList<HandleModel>) {
    use { ClassReader(it).accept(FindHandleClass(outHandle), 0) }
}

/**
 * `List<HandleModel>` 的扩展函数，将收集到的注解信息写入到 build 目录下的缓存文件中。
 *
 * 这个函数负责将内存中的 `HandleModel` 对象持久化到磁盘，以便后续的代码注入步骤可以读取。
 * 文件的命名基于 `HandleModel` 的子类名称（例如 `HandleModel.Autowired` -> `autowired.txt`）。
 *
 * @param outDir 缓存文件的输出目录，例如 `build/intermediates/router_scanning_data/debug`。
 */
private fun List<HandleModel>.write(outDir: Directory) {
    val cacheFileDir = outDir.asFile
    if (!cacheFileDir.exists()) cacheFileDir.mkdirs()
    // 按 HandleModel 的类型（Autowired, Module, Intercept, Initializer）进行分组
    this.groupBy { it::class.simpleName!! }.onEach { entry ->
        // 根据类型名生成文件名，例如 "autowired.txt"
        val fileName = "${entry.key.lowercase()}.txt"
        val file = outDir.file(fileName).asFile
        if (!file.exists()) file.createNewFile()
        val addBuffer = StringBuffer()
        // 遍历同类型的 HandleModel，将其信息序列化为字符串
        entry.value.forEachIndexed { index, model ->
            when (model) {
                is HandleModel.Autowired -> {
                    addBuffer.append(model.className)
                }

                is HandleModel.Module -> {
                    addBuffer.append(model.className)
                }

                is HandleModel.Initializer -> {
                    // 格式: "优先级 异步执行 类名"
                    addBuffer.append("${model.priority} ${model.async} ${model.className}")
                }

                is HandleModel.Intercept -> {
                    // 格式: "优先级 类名"
                    addBuffer.append("${model.priority} ${model.className}")
                }
            }
            if (index != entry.value.lastIndex) addBuffer.append("\n")
        }
        // 将序列化后的字符串写入文件
        file.writeText(addBuffer.toString())
    }

}